# -*- coding: utf-8 -*-
"""
等日间隔预测 2025 与 2026（默认：每4天、仅3/1~7/1窗口）
流程：
  1) 读取 P1sortdata.xlsx 并清洗，单位统一到百分比[0,100]
  2) 对历史每一年( <2025 )在窗口内做“回归后插值”，统一到等间隔栅格
  3) 按 DOY(年内天序)逐点拟合“Year→Value”的年际趋势（线性/二次/单调递增，可选鲁棒）
  4) 外推 2025、2026；采用“逐年递推的趋势缩放”抑制过冲
  5) 可选 Savitzky–Golay 平滑
输出：
  - output_p1/pred_2025_grid_every{RESOLUTION_DAYS}d.csv
  - output_p1/pred_2026_grid_every{RESOLUTION_DAYS}d.csv
  - output_p1/pred_2025_2026_all_every{RESOLUTION_DAYS}d.csv
"""

# =================== 参数区（可改） ===================
INPUT_XLSX = "P1sortdata.xlsx"
SHEET_NAME = 0

# 预测窗口与等间隔分辨率
WINDOW_START = (3, 1)         # 仅预测 3/1 ~ 7/1
WINDOW_END   = (7, 1)
RESOLUTION_DAYS = 4           # 等日间隔（天）

# 每年窗口内的“回归后插值”方式与阶数
REGRESSION_KIND = "poly"      # "poly" | "loess"
POLY_DEGREE_CANDIDATES = (1, 2, 3)
LOESS_FRAC = 0.6

# 逐 DOY 的年际趋势模型
YEAR_TREND = "linear"         # "linear" | "quadratic" | "isotonic"
ROBUST = "theilsen"           # "none" | "theilsen" | "ransac"  (仅对 linear 生效)
RIDGE_ALPHA = 0.0             # 线性/二次的L2正则（>0生效）

# 趋势缩放（防止外推过大）：逐年递推
# y_t = y_(t-1) + LAMBDA*(y_t_raw - y_(t-1))
TREND_DAMPING = 0.1           # 0<λ<=1，越小外推越保守

# 平滑（对最终序列）
SMOOTH = True
SMOOTH_WINDOW = 9             # 奇数，程序中会自动矫正
SMOOTH_POLY   = 2

OUTPUT_DIR = "./output_p1"
# =====================================================

import os, warnings, numpy as np, pandas as pd
from scipy.interpolate import interp1d
from scipy.signal import savgol_filter
warnings.filterwarnings("ignore")
os.makedirs(OUTPUT_DIR, exist_ok=True)

# ---------- I/O & 清洗 ----------
def _infer_cols(df: pd.DataFrame):
    date_candidates  = ["日期","date","时间","Date","time","Time"]
    value_candidates = ["就业率","就业率(%)","就业率（%）","y","value","rate"]
    dcol = None
    for c in df.columns:
        try: pd.to_datetime(df[c], errors="raise"); dcol = c; break
        except Exception: pass
    if dcol is None:
        for c in date_candidates:
            if c in df.columns: dcol = c; break
    vcol = None
    for c in df.columns:
        if c == dcol: continue
        s = pd.to_numeric(df[c], errors="coerce")
        if s.notna().sum() >= max(5, int(0.3*len(s))):
            vcol = c
            if any(k in str(c) for k in ["就","rate","value","%"]): break
    if vcol is None:
        for c in value_candidates:
            if c in df.columns: vcol = c; break
    if dcol is None or vcol is None:
        raise ValueError(f"无法识别日期/数值列：{df.columns.tolist()}")
    return dcol, vcol

def load_all(path, sheet=0):
    df = pd.read_excel(path, sheet_name=sheet)
    dcol, vcol = _infer_cols(df)
    df = df[[dcol, vcol]].copy()
    df.columns = ["date","y"]
    df["date"] = pd.to_datetime(df["date"], errors="coerce")
    df["y"]    = pd.to_numeric(df["y"], errors="coerce")
    df = df.dropna().drop_duplicates("date").sort_values("date").reset_index(drop=True)
    # 统一到百分比
    if df["y"].median() <= 1.5 and df["y"].max() <= 1.2:
        df["y"] *= 100.0
    df["y"] = df["y"].clip(0, 100)
    df["year"] = df["date"].dt.year
    return df

def make_grid(year, start_mmdd, end_mmdd, step_days):
    start = pd.Timestamp(year=year, month=start_mmdd[0], day=start_mmdd[1])
    end   = pd.Timestamp(year=year, month=end_mmdd[0],   day=end_mmdd[1])
    return pd.date_range(start, end, freq=f"{int(step_days)}D")

# ---------- 每年回归后插值（统一等间隔） ----------
try:
    from statsmodels.nonparametric.smoothers_lowess import lowess
except Exception:
    lowess = None
from sklearn.linear_model import Ridge
from sklearn.preprocessing import PolynomialFeatures
from sklearn.pipeline import Pipeline

def _aicc(n, sse, k):
    if n <= k + 1: return np.inf
    return n*np.log(max(sse/n, 1e-12)) + 2*k + (2*k*(k+1))/(n - k - 1)

def fit_poly_reg(x, y, deg, alpha=1e-6):
    mdl = Pipeline([
        ("poly", PolynomialFeatures(degree=deg, include_bias=True)),
        ("ridge", Ridge(alpha=alpha, fit_intercept=False, random_state=0))
    ])
    mdl.fit(x.reshape(-1,1), y)
    yhat = mdl.predict(x.reshape(-1,1))
    sse = float(np.sum((y - yhat)**2))
    return mdl, _aicc(len(y), sse, deg + 1)

def per_year_regression_interpolate(df_all, step_days, start_mmdd, end_mmdd, kind="poly"):
    years = sorted(df_all["year"].unique())
    target_grid = make_grid(2025, start_mmdd, end_mmdd, step_days)  # 用 2025 的 DOY 作为基准
    per_year_series = {}
    for y in years:
        g = df_all[df_all["year"] == y].copy()
        if g.empty: continue
        start = pd.Timestamp(year=y, month=start_mmdd[0], day=start_mmdd[1])
        end   = pd.Timestamp(year=y, month=end_mmdd[0],   day=end_mmdd[1])
        g_win = g[(g["date"] >= start) & (g["date"] <= end)].copy()
        g_use = g_win if len(g_win) >= 2 else g.copy()
        grid = make_grid(y, start_mmdd, end_mmdd, step_days)
        if len(g_use) < 2:
            if len(g) >= 2:
                xi = g["date"].astype("int64").values
                yi = g["y"].values
                try:   f = interp1d(xi, yi, kind="linear",  fill_value="extrapolate")
                except Exception: f = interp1d(xi, yi, kind="nearest", fill_value="extrapolate")
                yg = f(grid.astype("int64").values)
            else:
                yg = np.repeat(g["y"].iloc[0], len(grid))
            per_year_series[y] = pd.Series(np.clip(yg,0,100), index=grid); continue

        x = g_use["date"].dt.dayofyear.values.astype(float)
        yv = g_use["y"].values.astype(float)
        xg = pd.DatetimeIndex(grid).dayofyear.values.astype(float)

        if kind == "loess" and lowess is not None and len(g_use) >= 5:
            sm = lowess(yv, x, frac=LOESS_FRAC, it=0, return_sorted=True)
            f  = interp1d(sm[:,0], sm[:,1], kind="linear", fill_value="extrapolate")
            yg = f(xg)
        else:
            best = (None, np.inf, None)
            for deg in POLY_DEGREE_CANDIDATES:
                try:
                    m, aicc = fit_poly_reg(x, yv, deg)
                    if aicc < best[1]: best = (deg, aicc, m)
                except Exception:
                    continue
            model = best[2] if best[2] is not None else Ridge().fit(x.reshape(-1,1), yv)
            yg = model.predict(xg.reshape(-1,1))
        per_year_series[y] = pd.Series(np.clip(yg,0,100), index=grid)

    # 拼训练矩阵：行=2025的DOY网格；列=历史年份
    rows = []
    for dt in target_grid:
        doy = dt.timetuple().tm_yday
        row = {"date": dt}
        for y, ser in per_year_series.items():
            ser_doy = ser.index.dayofyear.values
            pos = np.where(ser_doy == doy)[0]
            j = int(pos[0]) if len(pos) else int(np.argmin(np.abs(ser_doy - doy)))
            row[y] = float(ser.iloc[j])
        rows.append(row)
    mat_df = pd.DataFrame(rows).set_index("date").sort_index()
    return mat_df

# ---------- 逐 DOY 年际趋势 ----------
from sklearn.linear_model import LinearRegression, TheilSenRegressor, RANSACRegressor
from sklearn.isotonic import IsotonicRegression

def fit_year_trend(x_years, y_values, year_trend="linear", robust="none", ridge_alpha=0.0):
    x = np.asarray(x_years, dtype=float).reshape(-1,1)
    y = np.asarray(y_values, dtype=float)
    mask = np.isfinite(y)
    if mask.sum() < 2: return None
    x, y = x[mask], y[mask]

    if year_trend == "linear":
        if robust == "theilsen":
            mdl = TheilSenRegressor(random_state=0).fit(x, y)
            return lambda yr: float(mdl.predict(np.array([[yr]], dtype=float)))
        elif robust == "ransac":
            base = LinearRegression()
            mdl = RANSACRegressor(base_estimator=base, random_state=0).fit(x, y)
            return lambda yr: float(mdl.predict(np.array([[yr]], dtype=float)))
        else:
            if ridge_alpha and ridge_alpha > 0:
                mdl = Ridge(alpha=ridge_alpha, fit_intercept=True, random_state=0).fit(x, y)
            else:
                mdl = LinearRegression().fit(x, y)
            return lambda yr: float(mdl.predict(np.array([[yr]], dtype=float)))

    elif year_trend == "quadratic":
        X = np.hstack([x, x**2])
        if ridge_alpha and ridge_alpha > 0:
            mdl = Ridge(alpha=ridge_alpha, fit_intercept=True, random_state=0).fit(X, y)
        else:
            mdl = LinearRegression().fit(X, y)
        return lambda yr: float(mdl.predict(np.array([[yr, yr**2]], dtype=float)))

    elif year_trend == "isotonic":
        order = np.argsort(x[:,0])
        iso = IsotonicRegression(increasing=True, out_of_bounds="clip").fit(x[order,0], y[order])
        return lambda yr: float(iso.predict([yr])[0])

    else:
        raise ValueError("YEAR_TREND 必须是 'linear' | 'quadratic' | 'isotonic'")

def smooth_series(y, window=9, poly=2):
    n = len(y)
    if n < 5: return y
    w = min(window, n - (1 - n % 2))
    if w < 3: w = 3
    if w % 2 == 0: w -= 1
    if poly >= w: poly = max(1, w - 1)
    try:
        return savgol_filter(y, window_length=w, polyorder=poly, mode="interp")
    except Exception:
        return y

# ---------- 预测主逻辑 ----------
def main():
    df_all = load_all(INPUT_XLSX, SHEET_NAME)

    # 历史年：全部 <2025 的年份
    df_hist = df_all[df_all["year"] < 2025].copy()
    if df_hist.empty:
        raise RuntimeError("无历史样本（<2025）。")

    # 每年回归后插值 → 训练矩阵(行=2025 DOY网格, 列=历史年份)
    mat_df = per_year_regression_interpolate(
        df_hist, RESOLUTION_DAYS, WINDOW_START, WINDOW_END, kind=REGRESSION_KIND
    )
    mat_df.to_csv(os.path.join(OUTPUT_DIR, f"train_matrix_every{RESOLUTION_DAYS}d.csv"),
                  encoding="utf-8-sig")

    years_hist = np.array(sorted([c for c in mat_df.columns if isinstance(c, (int, np.integer))]), dtype=float)
    grid2025 = mat_df.index
    grid2026 = make_grid(2026, WINDOW_START, WINDOW_END, RESOLUTION_DAYS)

    # 逐 DOY 拟合年际趋势 → 递推预测 2025、2026（带趋势缩放）
    y25_raw, y26_raw = [], []
    y25, y26 = [], []
    anchor_year = int(years_hist.max())  # 最近的历史年作为锚点

    for idx, row in mat_df.iterrows():
        vals = row[years_hist].values.astype(float)
        pred_fn = fit_year_trend(years_hist, vals, year_trend=YEAR_TREND, robust=ROBUST, ridge_alpha=RIDGE_ALPHA)
        if pred_fn is None:
            # 兜底：用非空均值
            m = float(np.nanmean(vals)) if np.isfinite(np.nanmean(vals)) else 0.0
            y25_raw.append(m); y26_raw.append(m)
            y25.append(m);    y26.append(m)
            continue

        # 原始外推
        y_anchor = float(pred_fn(anchor_year))
        y25_r = float(pred_fn(2025))
        y26_r = float(pred_fn(2026))
        y25_r = np.clip(y25_r, 0, 100)
        y26_r = np.clip(y26_r, 0, 100)

        # 递推缩放：先 2025 相对 anchor 缓推，再 2026 相对 2025 缓推
        y25_adj = y_anchor + TREND_DAMPING * (y25_r - y_anchor)
        y26_adj = y25_adj + TREND_DAMPING * (y26_r - y25_adj)

        y25_raw.append(y25_r); y26_raw.append(y26_r)
        y25.append(np.clip(y25_adj, 0, 100))
        y26.append(np.clip(y26_adj, 0, 100))

    y25 = np.array(y25); y26 = np.array(y26)

    # 平滑（可选）
    if SMOOTH:
        y25 = np.clip(smooth_series(y25, window=SMOOTH_WINDOW, poly=SMOOTH_POLY), 0, 100)
        y26 = np.clip(smooth_series(y26, window=SMOOTH_WINDOW, poly=SMOOTH_POLY), 0, 100)

    # 保存结果
    df25 = pd.DataFrame({"date": grid2025, "就业率（%）": y25, "year": 2025})
    df26 = pd.DataFrame({"date": grid2026, "就业率（%）": y26, "year": 2026})

    df25.to_csv(os.path.join(OUTPUT_DIR, f"pred_2025_grid_every{RESOLUTION_DAYS}d.csv"),
                index=False, encoding="utf-8-sig")
    df26.to_csv(os.path.join(OUTPUT_DIR, f"pred_2026_grid_every{RESOLUTION_DAYS}d.csv"),
                index=False, encoding="utf-8-sig")

    both = pd.concat([df25, df26], ignore_index=True)
    both.to_csv(os.path.join(OUTPUT_DIR, f"pred_2025_2026_all_every{RESOLUTION_DAYS}d.csv"),
                index=False, encoding="utf-8-sig")

    print("[Done] 预测完成。输出目录：", OUTPUT_DIR)
    print(f" - 2025：pred_2025_grid_every{RESOLUTION_DAYS}d.csv")
    print(f" - 2026：pred_2026_grid_every{RESOLUTION_DAYS}d.csv")
    print(f" - 合并：pred_2025_2026_all_every{RESOLUTION_DAYS}d.csv")
    print(f"\n关键参数：YEAR_TREND={YEAR_TREND}, ROBUST={ROBUST}, TREND_DAMPING={TREND_DAMPING}, SMOOTH={SMOOTH}")
    print("建议：若你感觉年际外推仍偏大，先把 TREND_DAMPING 调小（如 0.3）。")
    
if __name__ == "__main__":
    main()
